# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import logging
import json
from typing import Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.timestamp import get_timestamp
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.bitfinex_v2.auth import ApiKeyAuth


class BitfinexV2PrivateSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self, config, queue, worker_id, close_callback=None, logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._auth = ApiKeyAuth(config.api_key, config.api_secret)
    self._config = config
    self._channels = {}
    self._register_url('wss://api.bitfinex.com/ws/2')

    self._publisher = None
    self._init_publisher(queue, worker_id)
    self._ready = False

    # pong check
    self._last_pong_timestamp = None
    self._init_params = (config, queue, worker_id, close_callback, logger)

  def _check_pong_timestamp(self):
    if self._last_pong_timestamp is None:
      return

    time_elapsed = (get_timestamp() - self._last_pong_timestamp) / 1e9
    if time_elapsed > 90:
      self._logger.error('No pong received for %d seconds, restart trade client soon.',
                         time_elapsed)
      self.stop(trigger_close_callback=True)

  def _init_publisher(self, queue, worker_id: str):
    topic_private = topic_map.search_topic_id('bitfinex_v2_ws_private_%s' % worker_id)[0]

    self._publisher = {
        'private': queue.get_writer(topic_id=topic_private[1]),
    }

    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def _send_request(self, payload):
    if self._ws is not None:
      return self._ws.write_message(json.dumps(payload))
    else:
      self._logger.error('Subscriber %d: Websocket is closed!', id(self))
      return None

  @gen.coroutine
  def _ping(self):
    cid = 1000
    while not self.stopped:
      yield gen.sleep(30)
      cid += 1
      if self._ws is not None:
        self._logger.debug('ping')
        try:
          yield self._ws.write_message('{"event":"ping", "cid": %d}' % cid)
          self._check_pong_timestamp()
        except Exception:
          if self.stopped:
            self._logger.info('Bitfinex private subscriber (%d) is stopped already!', id(self))
          else:
            raise

  def request_position(self, symbol_list):
    # position_SYMBOL
    payload = [0, 'calc', None, [['position_%s' % symbol for symbol in symbol_list]]]
    self._send_request(payload)

  def request_wallet(self, wallet_type, symbol_list):
    # wallet_WALLET-TYPE_CURRENCY
    assert wallet_type in ('exchange', 'margin', 'funding')
    payload = [
        0, 'calc', None, [['wallet_%s_%s' % (wallet_type, symbol)] for symbol in symbol_list]
    ]
    self._send_request(payload)

  def place_order(self,
                  symbol: str,
                  side: None,
                  price: str,
                  qty: str,
                  client_order_id: int,
                  post_only=False,
                  hidden=False,
                  margin_trading=False):
    flags = 0
    if post_only:
      flags += 4096
    if hidden:
      flags += 64
    if margin_trading:
      order_type = 'LIMIT'
    else:
      order_type = 'EXCHANGE LIMIT'
    payload = [
        0,
        'on',
        None,
        {
            'cid': client_order_id,
            'type': order_type,
            'symbol': symbol,
            'amount': qty,
            'price': price,
            'flags': flags,
        }
    ]
    self._send_request(payload)

  def cancel_by_server_order_ids(self, server_order_ids: list):
    payload = [0, 'oc_multi', None, {'id': server_order_ids}]
    self._send_request(payload)

  @gen.coroutine
  def on_open(self):
    msg = yield self._ws.read_message()
    msg = json.loads(msg)
    if msg.get('event', None) != 'info' or msg.get('version', None) != 2:
      self._logger.error('Version mismatch: "%s". Stop subscription.' % msg)
      self._notify_close()
      return
    self._ioloop.add_callback(self._send_subscribe_request)
    self._ioloop.call_later(10, self._ping)

  def _generate_auth_payload(self):
    nonce = int(get_timestamp() / 1000)
    payload = 'AUTH%d' % nonce
    sig = self._auth.generate_signature(payload)
    return {
        'apiKey': self._config.api_key,
        'authSig': sig,
        'authNonce': str(nonce),
        'authPayload': payload,
        'event': 'auth'
    }

  @gen.coroutine
  def _send_subscribe_request(self):
    yield self._send_request(self._generate_auth_payload())
    """
    gen.sleep(1)
    payload = [
      0,
      "calc",
      None,
      [
        ["margin_sym_tBTCUSD"],
        ["position_tBTCUSD"],
        ["wallet_margin_BTC"],
        ["wallet_funding_USD"]
      ]
    ]
    yield self._ws.write_message(json.dumps(payload))
    """
    self._logger.info('Subscribe done')

  def on_message(self, timestamp: int, message: Union[dict, list]):
    if self.stopped:
      return

    topic = 'private'
    if type(message) == list:
      if not self._ready:
        self._logger.warn('Not ready')
        return
      channel = message[0]
      try:
        topic = self._channels[channel]
      except KeyError:
        self._logger.warning('Unknown channel: %s. Known channels are %s' %
                             (channel, self._channels))
        return
    elif type(message) == dict:
      if message['event'] == 'auth':
        channel = message['chanId']
        topic = 'private'
        self._channels[channel] = topic
        self._ready = True
        return
      elif message['event'] == 'info':
        pass
      elif message['event'] == 'pong':
        self._last_pong_timestamp = get_timestamp()
        self._logger.debug('%s, %s', message, id(self))
        return
      else:
        raise ValueError('Unknown message: %s' % message)
    try:
      self._publisher[topic].write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown topic: %s' % topic)

  def clone(self):
    return BitfinexV2PrivateSubscriber(*self._init_params)
